#include "rec.h"
#include "geometry_msgs/Pose.h"
#include "std_msgs/Bool.h"
#include <irb_msgs/ObjectPose.h>

#include <fstream>
using namespace std;

PoseInternal::PoseInternal()
{
	mX = 0;
	mY = 0;
	mZ = 0;
	mRx = 0;
	mRy = 0;
	mRz = 0;
	mRw = 0;
	mGripper = LIMBO;
}

PoseInternal::PoseInternal(double x, double y, double z,
						   double rx, double ry, double rz, double rw,
						   ros::Time timestamp)
{
	mX = x;
	mY = y;
	mZ = z;
	mRx = rx;
	mRy = ry;
	mRz = rz;
	mRw = rw;
	mTimestamp = timestamp;

	mGripper = LIMBO;
}

Recording::Recording(int id)
{
	objectID = id;
}

void Recording::HandlePose(const PoseInternal &pose)
{
	PoseInternal p(pose);
	// Calculate offset from first recorded pose
	if(mPoses.size())
		p.mOffset = p.mTimestamp - mPoses[0].mTimestamp;
	else
		p.mOffset = ros::Duration(0);

	mPoses.push_back(p);
}

void Recording::Save(string filename)
{
	ofstream out(filename.c_str());
	if(!out)
		throw "Couldn't open output file " + filename;

	out << "! " << objectID << endl;

	for(vector<PoseInternal>::iterator it = mPoses.begin(); it != mPoses.end(); it++)
	{
		PoseInternal pose = *it;
		out << "[[" << pose.mX << "," << pose.mY << "," << pose.mZ << "]";
		out << ",[" << pose.mRx << "," << pose.mRy << "," << pose.mRz << "," << pose.mRw << "]]";
		out << "!" << pose.mOffset.toNSec() << " " << pose.mGripper << endl;
	}

	out.close();
}

void Recording::Load(string filename)
{
	bool firstLine = true;

	ifstream in(filename.c_str());
	if(!in)
		throw "Couldn't open input file " + filename;

	mPoses.clear();

	while(in)
	{
		// Get a line of data
		char linebuf[1024];
		in.getline(linebuf, 1024);
		string str(linebuf);

		if(!in)
			break;

		// Remove [
		size_t pos = str.find('[');
		while(pos != string::npos)
		{
			str.replace(pos, 1, "");
			pos = str.find('[');
		}

		// Remove ]
		pos = str.find(']');
		while(pos != string::npos)
		{
			str.replace(pos, 1, "");
			pos = str.find(']');
		}

		// Remove ,
		pos = str.find(',');
		while(pos != string::npos)
		{
			str.replace(pos, 1, " ");
			pos = str.find(',');
		}

		// Remove !
		pos = str.find('!');
		while(pos != string::npos)
		{
			str.replace(pos, 1, " ");
			pos = str.find('!');
		}

		stringstream ss(str);

		// Load object ID
		if(firstLine)
		{
			ss >> objectID;
			firstLine = false;
		}
		// Load points (position and quatertion) + time offset + gripper state
		else
		{
			PoseInternal pose;
			long int nsec;

			ss >> pose.mX >> pose.mY >> pose.mZ >> pose.mRx >> pose.mRy >> pose.mRz >> pose.mRw;
			ss >> nsec >> pose.mGripper;
			pose.mOffset.fromNSec(nsec);
			pose.mTimestamp = ros::Time::now() + pose.mOffset;

			mPoses.push_back(pose);
		}
	}

	in.close();
}

RecordController::RecordController(ros::NodeHandle &node)
{
	mPlaybackPub = node.advertise<geometry_msgs::Pose>("playback", 1000);
	mPlaybackGripperPub = node.advertise<std_msgs::Bool>("gripper", 1000);
	mState = STANDBY;
	mPlaybackInterface = NULL;
}

void RecordController::HandlePose(const PoseInternal &pose, int objectID)
{
	// Discard if we are not recording
	if(mState != RECORDING)
		return;

	int recIndex = -1;
	for(size_t i = 0; i < mRecordings.size(); ++i)
		if(mRecordings[i]->ID() == objectID)
			recIndex = i;

	if(recIndex != -1)
		mRecordings[recIndex]->HandlePose(pose);
	else
	{
		Recording* newRec = new Recording(objectID);
		newRec->HandlePose(pose);
		mRecordings.push_back(newRec);
	}
}

void RecordController::Update(ros::Duration dt)
{
	switch(mState)
	{
		case STANDBY: break;
		case RECORDING: break;
		case PLAYBACK:
		{
			if(mPlaybackInterface)
			{
				const vector<PoseInternal> poses = mPlaybackInterface->Update(dt);
				for(vector<PoseInternal>::const_iterator it = poses.begin(); it != poses.end(); it++)
				{
					const PoseInternal& pose = *it;
					geometry_msgs::Pose msg;
					msg.position.x = pose.mX;
					msg.position.y = pose.mY;
					msg.position.z = pose.mZ;
					msg.orientation.x = pose.mRx;
					msg.orientation.y = pose.mRy;
					msg.orientation.z = pose.mRz;
					msg.orientation.w = pose.mRw;

					mPlaybackPub.publish(msg);

					// Sending gripper data
					std_msgs::Bool state;
					if(pose.mGripper == PoseInternal::CLOSED){
						state.data = false;
						mPlaybackGripperPub.publish(state);
					}
					else if(pose.mGripper == PoseInternal::OPEN){
						state.data = true;
						mPlaybackGripperPub.publish(state);
					}
				}

				if(mPlaybackInterface->AtEnd())
					mState = STANDBY;
			}
			break;
		}
	}
}

void RecordController::Play(Recording* rec)
{
	if(rec)
	{
		mState = PLAYBACK;
		if(mPlaybackInterface)
			delete mPlaybackInterface;
		mPlaybackInterface = new PlaybackInterface(rec);
	}
}

PlaybackInterface::PlaybackInterface(Recording *rec)
{
	mRecording = rec;
	mLastEmittedPose = -1; // no poses emitted yet
}

// Return a vector of poses to emit
vector<PoseInternal> PlaybackInterface::Update(ros::Duration dt)
{
	vector<PoseInternal> ret;

	// Increment the timer
	mPlaybackTimer += dt;

	// Get all poses from the current recording
	vector<PoseInternal>& poses = mRecording->Poses();

	// Start at the position after the one last emitted
	unsigned int poseIndex = mLastEmittedPose + 1;
	if(poseIndex < poses.size())
	{
		// Fetch the first pose
		PoseInternal p = poses[poseIndex];
		// Check if this much time has passed
		while(p.mOffset < mPlaybackTimer)
		{
			// If so, add it to the list
			ret.push_back(p);
			mLastEmittedPose = poseIndex;

			poseIndex++;
			// Stop if we are at the end of the list
			if(poseIndex >= poses.size())
				break;

			// Fetch the next pose
			p = poses[poseIndex];
		}
	}
	return ret;
}
